from util.util_image import load_image_from_path, load_image_from_bytes, calculate_ssim, evaluate_ssim
from ocr.ImageScoreOcrPaddle import ImageScoreOcrPaddle


# 检查是否有得分图片
def parse_images(list_image, target_image_data):
    ssim_list = []
    for image in list_image:
        image_data = load_image_from_bytes(image)
        ssim = calculate_ssim(image_data, target_image_data)
        # print("ssim:", ssim, evaluate_ssim(ssim))
        ssim_list.append(ssim)
    return ssim_list


class ScorerPart5:
    ANS_IMG_PATH_LIST_PART_5 = [f"files/ans/ans_5_1.png", f"files/ans/ans_5_2.png",
                                f"files/ans/ans_5_4.png"]
    ANS_IMG_OCR_KEYWORD_LIST = [[f"160.晶圆点胶", f"161.固化", f"162.加热", f"163.引脚键合", f"164.芯片注塑", f"165.激光打码"],
                                [f"激光打码设置菜单", f"输入您的姓名和学号"],
                                [f"你进行到了第166步", f"接下来可进行芯片测试虚拟仿真实验", f"BUPT"],
                                [f"BUPT", f"2018111999", f"北邮人"]]

    def __init__(self):
        self.paddleocr = ImageScoreOcrPaddle()
        self.ans_img_list_part_5 = []
        for path in self.ANS_IMG_PATH_LIST_PART_5:
            self.ans_img_list_part_5.append(load_image_from_path(path))

    # 检查是否有关键字
    def __ssim_contain_keyword(self, results):
        count_list = []
        for KEYWORD_LIST in self.ANS_IMG_OCR_KEYWORD_LIST:
            count = 0
            for result in results:
                if result["text"] in KEYWORD_LIST:
                    count += 1
            count_list.append(count / len(KEYWORD_LIST))
        print("count_list:", count_list, ", max:", max(count_list))
        return max(count_list)

    # 报告判分: 第五部分 -- SSIM
    def score_part_5_ssim(self, list_matched_images):
        score = 0
        ssim = 0
        print("[part 5][ssim]")
        # 检查是否有得分图片
        try:
            for target_img_data in self.ans_img_list_part_5:
                ssim_list = parse_images(list_matched_images, target_img_data)
                print("ssim_list:", ssim_list)
                # 有一张图相似度到阈值, 算得分
                if len(ssim_list) > 0 and max(ssim_list) > 0.5:
                    score = 1
                    ssim = max(ssim_list)
                    print("part 5: 得分, ssim_max:", max(ssim_list))
                    break
        except Exception as e:
            print(f"第5部分ssim判分出错: {e}")
        return score, ssim

    # 报告判分: 第五部分 -- OCR
    def score_part_5_ocr(self, list_matched_images):
        score = 0
        ssim = 0
        print("[part 5][ocr]")
        # 提取每张图的文字, 比较是否存在关键字
        try:
            for img_data in list_matched_images:
                # 提取每张图的文字
                results = self.paddleocr.get_ocr(img_data)
                if results is not None and len(results) > 0:
                    # 比较是否存在关键字
                    print("results:", results)
                    ssim = self.__ssim_contain_keyword(results)
                    if ssim >= 0.5:
                        print("part 5: ocr得分, ssim_max:", ssim)
                        score = 1
                        break
        except Exception as e:
            print(f"第5部分orc判分出错: {e}")
        return score, ssim

    # 报告判分: 第五部分
    def score_part_5(self, dict_matched_paragraphs, list_matched_images):
        score, ssim = self.score_part_5_ssim(list_matched_images)
        if score == 0:
            score, ssim = self.score_part_5_ocr(list_matched_images)
        return score, ssim